package my.tree;

import java.util.Comparator;

/**
 * <p>
 * AVL 树: 平衡的二叉搜索树
 * </p>
 *
 * @author AJun
 * @since 2020年8月20日
 */
@SuppressWarnings("ALL")
public class AVLTree<E> extends BBST<E> {

    /**
     * 空参构造
     * <p>
     * 没有传入比较器，则存储的元素必须 implements Comparable 接口实现比较逻辑，否则会报异常
     * </p>
     */
    public AVLTree() {
        this(null);
    }

    /**
     * 有参构造
     *
     * @param comparator 用于元素比较的比较器
     */
    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    /**
     * 添加之后的平衡调整
     *
     * @param node 添加的节点
     */
    @Override
    protected void afterAdd(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                // 此节点平衡，更新高度
                updateHeight(node);
            } else {
                // 恢复平衡
                rebalance(node);
                // 该节点恢复平衡后，整棵树恢复平衡
                return;
            }
        }
    }

    /**
     * 删除之后的平衡调整
     *
     * @param node 被删除的节点 或者 用于取代被删除节点的子节点(当被删除节点度为1时)
     */
    @Override
    protected void afterRemove(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                // 此节点平衡，更新高度
                updateHeight(node);
            } else {
                // 恢复平衡，之后再寻找父节点看是否平衡，如果不平衡需要继续恢复平衡
                rebalance(node);
            }
        }
    }

    /**
     * 旋转后的节点维护
     *
     * @param grand  祖父节点
     * @param parent 父节点
     * @param child  父节点的左子节点或右子节点
     */
    @Override
    protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
        // 旋转后的节点维护完毕后，更新高度
        super.afterRotate(grand, parent, child);

        // 更新 grand 和 parent 的高度
        updateHeight(grand);
        updateHeight(parent);
    }

    /**
     * 恢复平衡
     *
     * @param grand 高度最低的不平衡节点
     */
    private void rebalance(Node<E> grand) {
        Node<E> parent = ((AVLNode<E>) grand).tallerChild();
        Node<E> node = ((AVLNode<E>) parent).tallerChild();
        if (parent.isLeftChild()) { // L
            if (node.isLeftChild()) { // LL
                rotateRight(grand);
            } else { // LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        } else { // R
            if (node.isLeftChild()) { // RL
                rotateRight(parent);
                rotateLeft(grand);
            } else { // RR
                rotateLeft(grand);
            }
        }
    }

    /**
     * 是否平衡，平衡因子的绝对值小于等于 1
     *
     * @return true: 平衡; false: 不平衡
     */
    private boolean isBalanced(Node<E> node) {
        return Math.abs(((AVLNode<E>) node).balanceFactor()) <= 1;
    }

    /**
     * 更新节点高度
     *
     * @param node 要更新高度的节点
     */
    private void updateHeight(Node<E> node) {
        ((AVLNode<E>) node).updateHeight();
    }

    /**
     * 创建AVL树节点
     *
     * @param element 节点的元素内容
     * @param parent  节点的父节点
     * @return AVL树节点
     */
    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }

    /**
     * AVL 树节点
     */
    private static class AVLNode<E> extends Node<E> {
        /**
         * 节点高度。默认设为 1，因为新加的节点都是叶子节点，高度为1
         */
        int height = 1;

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /**
         * 平衡因子计算
         *
         * @return 平衡因子
         */
        public int balanceFactor() {
            int leftHeight = (left == null) ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = (right == null) ? 0 : ((AVLNode<E>) right).height;
            return leftHeight - rightHeight;
        }

        /**
         * 更新节点高度
         */
        public void updateHeight() {
            int leftHeight = (left == null) ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = (right == null) ? 0 : ((AVLNode<E>) right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        /**
         * 获取该节点的左右子树高度较高的子树
         *
         * @return 该节点的左右子树高度较高的子树
         */
        public Node<E> tallerChild() {
            int leftHeight = (left == null) ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = (right == null) ? 0 : ((AVLNode<E>) right).height;
            if (leftHeight > rightHeight) {
                return left;
            }
            if (leftHeight < rightHeight) {
                return right;
            }
            // 如果相等则返回与此节点同方向的节点
            return isLeftChild() ? left : right;
        }

        @Override
        public String toString() {
            String parentStr = "null";
            if (parent != null) {
                parentStr = parent.element.toString();
            }
            return element + "_p(" + parentStr + ")_h(" + height + ")";
        }
    }

}
